class Solution(object):
    def findMinHeightTrees(self, n, edges):
        """
        :type n: int
        :type edges: List[List[int]]
        :rtype: List[int]
        """
        if n==1: return [0]
        neighbours = [set() for _ in range(n)]
        for p1,p2 in edges:
            neighbours[p1].add(p2)
            neighbours[p2].add(p1)
        
        leaves = [i for i in range(n) if len(neighbours[i])==1]
        while n>2:
            n -= len(leaves)
            newLeaves = []
            for i in leaves:
                j = neighbours[i].pop()
                neighbours[j].remove(i)
                if len(neighbours[j])==1: newLeaves.append(j)
            leaves = newLeaves
        return leaves

print(Solution().findMinHeightTrees(4, [[1, 0], [1, 2], [1, 3]]))